package client

import (
	"bytes"
	"crypto/tls"
	"encoding/binary"
	"fmt"

	"github.com/pingcap/errors"
	. "server/dbtools/xcanal/mysql"
	"server/dbtools/xcanal/packet"
)

const defaultAuthPluginName = AUTH_NATIVE_PASSWORD

var supportedAuthPlugins = []string{AUTH_NATIVE_PASSWORD, AUTH_SHA256_PASSWORD, AUTH_CACHING_SHA2_PASSWORD}

func authPluginAllowed(pluginName string) bool {
	for _, p := range supportedAuthPlugins {
		if pluginName == p {
			return true
		}
	}
	return false
}

func (c *Conn) readInitialHandshake() error {
	data, err := c.ReadPacket()
	if err != nil {
		return errors.Trace(err)
	}

	if data[0] == ERR_HEADER {
		return errors.Annotate(c.handleErrorPacket(data), "read initial handshake error")
	}

	if data[0] < MinProtocolVersion {
		return errors.Errorf("invalid protocol version %d, must >= 10", data[0])
	}

	pos := 1 + bytes.IndexByte(data[1:], 0x00) + 1

	c.connectionID = binary.LittleEndian.Uint32(data[pos : pos+4])
	pos += 4

	c.salt = []byte{}
	c.salt = append(c.salt, data[pos:pos+8]...)

	pos += 8 + 1

	c.capability = uint32(binary.LittleEndian.Uint16(data[pos : pos+2]))
	if c.capability&CLIENT_PROTOCOL_41 == 0 {
		return errors.New("the MySQL server can not support protocol 41 and above required by the client")
	}
	if c.capability&CLIENT_SSL == 0 && c.tlsConfig != nil {
		return errors.New("the MySQL Server does not support TLS required by the client")
	}
	pos += 2

	if len(data) > pos {
		pos += 1

		c.status = binary.LittleEndian.Uint16(data[pos : pos+2])
		pos += 2
		c.capability = uint32(binary.LittleEndian.Uint16(data[pos:pos+2]))<<16 | c.capability
		pos += 2

		maxAuthDataLen := 21
		if c.capability&CLIENT_PLUGIN_AUTH != 0 && int(data[pos]) > maxAuthDataLen {
			maxAuthDataLen = int(data[pos])
		}

		pos += 10 + 1

		resetOfAuthDataEndPos := pos + maxAuthDataLen - 8 - 1
		c.salt = append(c.salt, data[pos:resetOfAuthDataEndPos]...)

		pos = resetOfAuthDataEndPos + 1

		if c.capability&CLIENT_PLUGIN_AUTH != 0 {
			c.authPluginName = string(data[pos : len(data)-1])
		}
	}

	if c.authPluginName == "" {
		c.authPluginName = defaultAuthPluginName
	}

	return nil
}

func (c *Conn) genAuthResponse(authData []byte) ([]byte, bool, error) {
	switch c.authPluginName {
	case AUTH_NATIVE_PASSWORD:
		return CalcPassword(authData[:20], []byte(c.password)), false, nil
	case AUTH_CACHING_SHA2_PASSWORD:
		return CalcCachingSha2Password(authData, c.password), false, nil
	case AUTH_CLEAR_PASSWORD:
		return []byte(c.password), true, nil
	case AUTH_SHA256_PASSWORD:
		if len(c.password) == 0 {
			return nil, true, nil
		}
		if c.tlsConfig != nil || c.proto == "unix" {
			return []byte(c.password), true, nil
		} else {
			return []byte{1}, false, nil
		}
	default:
		return nil, false, fmt.Errorf("auth plugin '%s' is not supported", c.authPluginName)
	}
}

func (c *Conn) genAttributes() []byte {
	if len(c.attributes) == 0 {
		return nil
	}

	attrData := make([]byte, 0)
	for k, v := range c.attributes {
		attrData = append(attrData, PutLengthEncodedString([]byte(k))...)
		attrData = append(attrData, PutLengthEncodedString([]byte(v))...)
	}
	return append(PutLengthEncodedInt(uint64(len(attrData))), attrData...)
}

func (c *Conn) writeAuthHandshake() error {
	if !authPluginAllowed(c.authPluginName) {
		return fmt.Errorf("unknow auth plugin name '%s'", c.authPluginName)
	}

	capability := CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION |
		CLIENT_LONG_PASSWORD | CLIENT_TRANSACTIONS | CLIENT_PLUGIN_AUTH
	capability |= c.capability & CLIENT_LONG_FLAG
	capability |= c.ccaps&CLIENT_FOUND_ROWS | c.ccaps&CLIENT_IGNORE_SPACE |
		c.ccaps&CLIENT_MULTI_STATEMENTS | c.ccaps&CLIENT_MULTI_RESULTS |
		c.ccaps&CLIENT_PS_MULTI_RESULTS | c.ccaps&CLIENT_CONNECT_ATTRS

	if c.tlsConfig != nil {
		capability |= CLIENT_SSL
	}

	auth, addNull, err := c.genAuthResponse(c.salt)
	if err != nil {
		return err
	}

	var authRespLEIBuf [9]byte
	authRespLEI := AppendLengthEncodedInteger(authRespLEIBuf[:0], uint64(len(auth)))
	if len(authRespLEI) > 1 {
		capability |= CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA
	}

	length := 4 + 4 + 1 + 23 + len(c.user) + 1 + len(authRespLEI) + len(auth) + 21 + 1
	if addNull {
		length++
	}
	if len(c.db) > 0 {
		capability |= CLIENT_CONNECT_WITH_DB
		length += len(c.db) + 1
	}
	attrData := c.genAttributes()
	if len(attrData) > 0 {
		capability |= CLIENT_CONNECT_ATTRS
		length += len(attrData)
	}

	data := make([]byte, length+4)

	data[4] = byte(capability)
	data[5] = byte(capability >> 8)
	data[6] = byte(capability >> 16)
	data[7] = byte(capability >> 24)

	data[8] = 0x00
	data[9] = 0x00
	data[10] = 0x00
	data[11] = 0x00

	data[12] = DEFAULT_COLLATION_ID

	if c.tlsConfig != nil {
		if err := c.WritePacket(data[:(4+4+1+23)+4]); err != nil {
			return err
		}

		tlsConn := tls.Client(c.Conn.Conn, c.tlsConfig)
		if err := tlsConn.Handshake(); err != nil {
			return err
		}

		currentSequence := c.Sequence
		c.Conn = packet.NewConn(tlsConn)
		c.Sequence = currentSequence
	}

	pos := 13
	for ; pos < 13+23; pos++ {
		data[pos] = 0
	}

	if len(c.user) > 0 {
		pos += copy(data[pos:], c.user)
	}
	data[pos] = 0x00
	pos++

	pos += copy(data[pos:], authRespLEI)
	pos += copy(data[pos:], auth)
	if addNull {
		data[pos] = 0x00
		pos++
	}

	if len(c.db) > 0 {
		pos += copy(data[pos:], c.db)
		data[pos] = 0x00
		pos++
	}

	pos += copy(data[pos:], c.authPluginName)
	data[pos] = 0x00
	pos++

	if len(attrData) > 0 {
		copy(data[pos:], attrData)
	}

	return c.WritePacket(data)
}
